package com.demo.elastic;

import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
import org.apache.commons.lang.StringUtils;
import org.elasticsearch.action.admin.indices.create.CreateIndexResponse;
import org.elasticsearch.action.admin.indices.delete.DeleteIndexResponse;
import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.cluster.metadata.MappingMetaData;
import org.elasticsearch.common.collect.ImmutableOpenMap;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

public class ElasticExporterApplication {

    private static final Logger logger = LoggerFactory.getLogger(ElasticExporterApplication.class);

    private Client sourceClient;

    private Client targetClient;

    private ExporterConfiguration configuration;

    private ElasticExporterApplication(String[] args) {
        configuration = ExporterConfiguration.parse(args);
    }

    public static void main(String[] args) {
        new ElasticExporterApplication(args).postConstruct();
    }


    private void postConstruct() {
        Runtime.getRuntime().addShutdownHook(new Thread(() -> configuration.destroy()));
        this.sourceClient = configuration.sourceClient();
        this.targetClient = configuration.targetClient();
        for (ObjectObjectCursor<String, ImmutableOpenMap<String, MappingMetaData>> index : getSourceIndices()) {
            if (StringUtils.startsWith(index.key, "zipkin:")) {
                logger.info("skip index::{}", index.key);
                continue;
            }
            recreateTargetIndex(index.key);
            for (ObjectObjectCursor<String, MappingMetaData> type : index.value) {
                export(index.key, type.key);
            }
        }
    }

    private void bulkIndexToTarget(String index, String type, SearchHit[] datas) {
        if (datas.length == 0) {
            return;
        }
        BulkRequestBuilder bulkRequest = targetClient.prepareBulk();

        for (SearchHit data : datas) {
            IndexRequestBuilder builder = targetClient.prepareIndex(index, type, data.getId()).setSource(data.source());
            bulkRequest.add(builder.request());
        }

        BulkResponse bulkResponse = bulkRequest.execute().actionGet();
        int errorCount = 0;
        if (bulkResponse.hasFailures()) {
            Map<String, String> failedDocuments = new HashMap<>();
            for (BulkItemResponse item : bulkResponse.getItems()) {
                if (item.isFailed()) {
                    failedDocuments.put(item.getId(), item.getFailureMessage());
                    errorCount++;
                }
            }
            System.out.println("Bulk indexing has failures. Use ElasticsearchException.getFailedDocuments() for detailed messages ["
                    + failedDocuments + "]");
        }

        logger.info("sent {} objects to destination elasticsearch[{}::{}], wrote::{}", datas.length,
                index, type, bulkResponse.getItems().length - errorCount);
    }


    private void export(String index, String type) {
        final int size = configuration.getSize();
        SearchRequestBuilder search = sourceClient.prepareSearch(index).setTypes(type);
        search.setSize(size);
        search.setQuery(QueryBuilders.queryStringQuery("*:*"));
        search.setScroll(TimeValue.timeValueMinutes(2));
        SearchResponse scrollResp = search.get();
        int offset = 0;
        if (scrollResp.getHits().getTotalHits() > 0) {
            do {
                logger.info("got {} objects from source elasticsearch[{}::{}] (offset::{}, total::{})",
                        scrollResp.getHits().getHits().length,
                        index, type, offset * size, scrollResp.getHits().getTotalHits());
                bulkIndexToTarget(index, type, scrollResp.getHits().getHits());
                scrollResp = sourceClient.prepareSearchScroll(scrollResp.getScrollId())
                        .setScroll(TimeValue.timeValueMinutes(2)).execute().actionGet();
                offset++;
            } while (scrollResp.getHits().getHits().length != 0);
        }
    }

    private void recreateTargetIndex(String index) {
        try {
            DeleteIndexResponse response = targetClient.admin().indices().prepareDelete(index).execute().actionGet();
            logger.info("delete target index [{}] result::{}", index, response.isAcknowledged());
        } catch (Exception ex) {
            logger.warn("delete target index [{}] message::{}", index, ex.getMessage());
        }
        CreateIndexResponse createResponse = targetClient.admin().indices().prepareCreate(index).execute().actionGet();
        logger.info("recreate target index [{}] result::{}", index, createResponse.isAcknowledged());
    }

    private ImmutableOpenMap<String, ImmutableOpenMap<String, MappingMetaData>> getSourceIndices() {

        GetIndexRequest request = new GetIndexRequest()
                .features(GetIndexRequest.Feature.MAPPINGS);
        return sourceClient.admin().indices().getIndex(request).actionGet().getMappings();
    }


}
